import numpy

from physics.vector.array import _array 

class Vector3Array(object):
    '''
    NumPy-like array of 3-vectors.
    
    Arguments
    ----------
    w : sequence of 3-tuples of length N
        A sequence of (x,y,z) sequences to be stored in the array

    OR 

    x : sequence of length N
        The x-components of the vectors to be stored in the array
    y : sequence of length N
        The y-components of the vectors to be stored in the array
    z : sequence of length N
        The z-components of the vectors to be stored in the array
    
    Keyword Arguments
    -------------------
    dtype: NumPy datatype object
        The data type used for storage of the vector components (default is `numpy.float64')
    nan_to_num: bool
        if True, all functions will return arrays filtered through NumPy `nan_to_num' function to 
        remove nans and infs.
    length: int
        Accept three scalar arguments and create Vector3Array of length `length'
    '''

    def __init__(self, *args, **kwargs):
        # TODO:
        # - add support for setting r, theta, phi
        # - add support for setting pt, eta, phi
        # - add support for setting pt, theta, phi

        if 'dtype' in kwargs:
            self.dtype = kwargs['dtype']
        else:
            self.dtype = numpy.float64

        if 'nan_to_num' in kwargs:
            self.nan_to_num = kwargs['nan_to_num']
        else:
            self.nan_to_num = False

        if 'length' in kwargs:
            self.shape = (kwargs['length'],)
        else:
            self.shape = None

        if len(args) == 1:
            if sum(map(len, args[0])) != 3*len(args[0]):
                msg = "Calling Vector3Array.__init__() with a single argument requires that argument to be"
                msg += " a sequence of 3-tuples"
                raise ValueError(msg)
    
            x, y, z = zip(*args[0])

        elif len(args) == 3:
            x, y, z = args
        
        else:
            msg = "Calling Vector3Array.__init__() with {0} arguments is undefined".format(len(args))
            raise ValueError(msg)
 
        if numpy.isscalar(x):            
            if numpy.isscalar(y):
                if numpy.isscalar(z):
                    if self.shape is None:
                        msg = "Cannot initialize Vector3Array with 3 scalars without `length' kwarg".format(self.x.shape)
                        raise ValueError(msg)
                else:
                    self.shape = (len(z),)
            else:
                self.shape = (len(y),)                
        else:
            self.shape = (len(x),)
                                
        if numpy.isscalar(x):
            self.x = numpy.ones(self.shape, dtype=self.dtype) * x
        else:            
            self.x = numpy.array(x, dtype=self.dtype)

        if numpy.isscalar(y):
            self.y = numpy.ones(self.shape, dtype=self.dtype) * y
        else:            
            self.y = numpy.array(y, dtype=self.dtype)

        if numpy.isscalar(z):
            self.z = numpy.ones(self.shape, dtype=self.dtype) * z
        else:            
            self.z = numpy.array(z, dtype=self.dtype)
                
        if len(self.x.shape) > 1:
            msg = "Argument 1 has shape {0}; must be a one-dimensional sequence".format(self.x.shape)
            raise ValueError(msg)

        if len(self.y.shape) > 1:
            msg = "Argument 2 has shape {0}; must be a one-dimensional sequence".format(self.y.shape)
            raise ValueError(msg)

        if len(self.z.shape) > 1:
            msg = "Argument 3 has shape {0}; must be a one-dimensional sequence".format(self.z.shape)
            raise ValueError(msg)

        lengths = dict(zip('xyz', map(len, [self.x, self.y, self.z])))
        
        if lengths['x'] != lengths['y'] or lengths['y'] != lengths['z'] or lengths['x'] != lengths['z']:
            msg = "Arguments do not all have the same length: {0}".format(lengths)
            raise ValueError(msg)

        
        self.shape = self.x.shape

    def phi(self):
        '''
        Returns an array containing the azimuthal angle on (-pi, pi) for each vector.
        '''
        
        if self.nan_to_num:
            return numpy.nan_to_num(numpy.arctan2(self.y, self.x))

        return numpy.arctan2(self.y, self.x)
        
    def theta(self):
        '''
        Returns an array containing the polar angle on (0, pi) for each vector.
        '''

        if self.nan_to_num:
            return numpy.nan_to_num(numpy.arctan2(self.pt(), self.z))

        return numpy.arctan2(self.pt(), self.z)
        
    def cos_theta(self):
        '''
        Returns an array containing the cosine of the polar angle for each vector.
        '''

        if self.nan_to_num:
            return numpy.nan_to_num(_array.v3_cos_theta(self.x, self.y, self.z))
        
        return _array.v3_cos_theta(self.x, self.y, self.z)


    def mag2(self):
        '''
        Returns an array containing the magnitude squared of each vector.
        '''

        return self.x**2 + self.y**2 + self.z**2

    def mag(self):
        '''
        Returns an array containing the magnitude of each vector.
        '''

        return numpy.sqrt(self.mag2())

    def pt2(self, axis=None):
        '''
        Returns an array containing the square of the transverse component of each vector.

        Optional Parameters:
        --------------------
        axis: Vector3Array of length N or sequence of length 3
            The axes or axis relative to which the transverse component should be calculated.
        '''

        if axis:
            if type(axis) == type(self):
                if axis.shape != self.shape:
                    raise ValueError("Axis shape {0} doesn't match data shape {1}".format(axis.shape, data.shape))
                return _array.v3_pt2(self.x, self.y, self.z, axis.x, axis.y, axis.z)
            elif len(axis) == 3:
                return _array.v3_pt2_single_axis(self.x, self.y, self.z, 
                                                 float(axis[0]), float(axis[1]), float(axis[2]))

            else:
                msg = "`axis' argument must be a Vector3Array of shape `{0}' or a 3-tuple".format(self.shape)
                raise ValueError(msg)

        else:
            return self.x**2 + self.y**2

    def pt(self, axis=None):
        '''
        Returns an array containing the transverse component of each vector.

        Optional Parameters:
        --------------------
        axis: Vector3Array of length N or sequence of length 3
            The axes or axis relative to which the transverse component should be calculated.
        '''

        return numpy.sqrt(self.pt2(axis))


    def __add__(self, rhs):

        if type(rhs) != type(self):
            raise ValueError("Cannot add type `{0}' to type `{1}'".format(type(rhs),
                                                                          type(self)))

        return Vector3Array(self.x + rhs.x, 
                       self.y + rhs.y,
                       self.z + rhs.z)


    def __sub__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot subtract type `{0}' from type `{1}'".format(type(rhs),
                                                                                 type(self)))
        return Vector3Array(self.x - rhs.x, 
                       self.y - rhs.y,
                       self.z - rhs.z)
        

    def __mul__(self, rhs):
        if numpy.isscalar(rhs):
            return Vector3Array(rhs*self.x,
                           rhs*self.y,
                           rhs*self.z)
        elif type(rhs) == type(self):
            return self.dot(rhs)
    
        else:
            msg = "Multiplication undefined between types `{0}' and `{1}'".format(type(rhs),
                                                                                  type(self))
            raise ValueError(msg)


    def __div__(self, rhs):
        if numpy.isscalar(rhs):
            return Vector3Array(self.x/rhs,
                           self.y/rhs,
                           self.z/rhs)
        else:
            msg = "Can only scale with a scalar; input is of type `{1}'".format(type(rhs))
            raise ValueError(msg)
    


    def __iadd__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot add type `{0}' to type `{1}'".format(type(rhs),
                                                                          type(self)))
        self.x += rhs.x
        self.x += rhs.y
        self.x += rhs.z

        return self

    def __isub__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot subtract type `{0}' from type `{1}'".format(type(rhs),
                                                                                 type(self)))
        self.x -= rhs.x
        self.x -= rhs.y
        self.x -= rhs.z

        return self

    def __imul__(self, rhs):
        if numpy.isscalar(rhs) or hasattr(rhs, 'shape') and len(rhs.shape) == 1:
            self.x *= rhs
            self.y *= rhs
            self.z *= rhs
        else:
            msg = "Can only scale with a scalar or scalar array; input is of type `{1}'".format(type(rhs))
            raise ValueError(msg)

        return self

    def __idiv__(self, rhs):
        if numpy.isscalar(rhs):
            self.x /= rhs
            self.y /= rhs
            self.z /= rhs
        else:
            msg = "Can only scale with a scalar; input is of type `{1}'".format(type(rhs))
            raise ValueError(msg)

        return self

    def __getitem__(self, index):
        return self.x[index], self.y[index], self.z[index]

    def __len__(self):
        return self.shape[0]
            
    def __eq__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot compare type `{0}' to type `{1}'".format(type(rhs),
                                                                              type(self)))

        return self.x == rhs.x and self.y == rhs.y and self.z == rhs.z

    def __ne__(self, rhs):
        if type(rhs) != type(self):
            raise ValueError("Cannot compare type `{0}' to type `{1}'".format(type(rhs),
                                                                              type(self)))

        return self.x != rhs.x and self.y != rhs.y and self.z != rhs.z

    def __neg__(self):
        return Vector3Array(-self.x, -self.y, -self.z)

    def __str__(self):
        return str(numpy.array(zip(self.x, self.y, self.z)))

    def __iter__(self):
        for i in xrange(self.shape[0]):
            yield (self.x[i], self.y[i], self.z[i])            
    
#    def next(self):

    def dot(self, rhs):        
        return self.x*rhs.x + self.y*rhs.y + self.z*rhs.z

    def cross(self, rhs):
        return Vector3Array(self.y*rhs.z - self.z*rhs.y, 
                       self.z*rhs.x - self.x*rhs.z,
                       self.x*rhs.y - self.y*rhs.x)

    def unit(self):
        uv = _array.v3_unit(self.x, self.y, self.z)
        return Vector3Array(uv[0], uv[1], uv[2])    

    def angle(self, rhs):
        '''
        Returns an array containing the angle of each vector with respect to the vectors in `rhs'.

        Parameters:
        --------------------
        rhs: Vector3Array of length N
            The vectors with respect to which the angles should be calculated.
        '''   

        if type(rhs) != type(self):
            raise ValueError("Cannot find angle between type `{0}' and type `{1}'".format(type(rhs),
                                                                                          type(self)))
    
        return _array.v3_angle(self.x, self.y, self.z, 
                               rhs.x, rhs.y, rhs.z)


    def eta(self):
        '''
        Returns an array containing the pseudorapidity for each vector.
        '''
        
        return _array.v3_eta(self.x, self.y, self.z)
    

    def rotate_x(self, angle):
        '''
        Rotates each vector in place by angle `angle' around the X axis.

        Parameters:
        --------------------
        angle: float
            The angle of rotation about the X axis
        '''   

        if not numpy.isscalar(angle):
            raise ValueError("Argument 1 must be a scalar")
        
        _array.v3_rotate_x(self.x, self.y, self.z, angle)

    
    def rotate_y(self, angle):
        '''
        Rotates each vector *in place* by angle `angle' around the Y axis.

        Parameters:
        --------------------
        angle: float
            The angle of rotation about the Y axis
        '''   

        if not numpy.isscalar(angle):
            raise ValueError("Argument 1 must be a scalar")
        
        _array.v3_rotate_y(self.x, self.y, self.z, angle)


    def rotate_z(self, angle):
        '''
        Rotates each vector *in place* by angle `angle' around the Z axis.

        Parameters:
        --------------------
        angle: float
            The angle of rotation about the Z axis
        '''   

        if not numpy.isscalar(angle):
            raise ValueError("Argument 1 must be a scalar")
        
        _array.v3_rotate_z(self.x, self.y, self.z, angle)    
        

        

